Skip to content

Preserve element type for mvcombine ARRAY_AGG on the DataFusion route - #22617

Open
noCharger wants to merge 1 commit into
opensearch-project:mainfrom
noCharger:mvcombine-datafusion-arrayagg
Open

Preserve element type for mvcombine ARRAY_AGG on the DataFusion route#22617
noCharger wants to merge 1 commit into
opensearch-project:mainfrom
noCharger:mvcombine-datafusion-arrayagg

Conversation

@noCharger

@noCharger noCharger commented Jul 31, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Description

PPL mvcombine lowers to Calcite ARRAY_AGG. On the analytics-engine (DataFusion) route this currently fails with No enum constant org.opensearch.analytics.spi.AggregateFunction.ARRAY_AGG — the aggregate is not recognized. Mapping it onto the existing list()/values() array-agg operator is not correct either: that operator casts every element to VARCHAR by design (the PPL list()/values() multivalue-string contract), so mvcombine over a non-string field would return ARRAY<VARCHAR> instead of preserving the source type.

This adds a dedicated type-preserving path so the collected array keeps the source element type (e.g. ARRAY<INT>, not ARRAY<VARCHAR>):

  • AggregateFunction.fromNameOrError — recognize ARRAY_AGG as the LIST (collect-into-array) family for backend viability and the STATE_EXPANDING decomposition. Element typing is handled separately in the rewriter, so this mapping only affects viability/decomposition shape.
  • DataFusionFragmentConvertor — add LOCAL_ARRAY_AGG_TYPED_OP, a LocalAggOp whose return type is derived from the operand (not hardcoded VARCHAR) and which does not cast the element. It drops null elements (matching ARRAY_AGG ... FILTER (IS NOT NULL)) via filtersNullArgs, and maps to the same type-generic DataFusion array_agg UDAF as LOCAL_ARRAY_AGG_OP.
  • PplAggregateCallRewriter — route ARRAY_AGG to the typed op for the PARTIAL/SINGLE stage, and to the existing (already type-generic) list_merge for the FINAL merge.

list()/values() are unchanged (still ARRAY<VARCHAR>). No Rust change — the array_agg/list_merge UDAFs already derive the element type from the input.

Verified on the DataFusion route (composite/parquet index, analytics.planner.prefer_metadata_driver=false): source=<idx> | fields category, age | mvcombine age over an integer field returns unquoted integers ([30, 30], [25, 35], [40]) with the query profile reporting chosen_backend=datafusion; list()/values() still return ARRAY<VARCHAR>.

Related Issues

Enables mvcombine on the analytics-engine (DataFusion) route. The mvcombine command lives in the SQL plugin (opensearch-project/sql#4766).

Check List

  • Functionality includes testing.
    • AggregateFunctionTests#testArrayAggResolvesToList — ARRAY_AGG resolves to the LIST family.
    • DataFusionFragmentConvertorTests#testAttachFragmentOnTop_ArrayAggOverScalar_MeasureArgIsRawFieldNotVarcharCast — the substrait measure arg is the raw field, not a VARCHAR cast (the type-preservation guarantee).
  • Commits are signed per the DCO using --signoff.

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.

@noCharger
noCharger requested a review from a team as a code owner July 31, 2026 04:07
@github-actions

github-actions Bot commented Jul 31, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

PR Reviewer Guide 🔍

(Review updated until commit f6cfd03)

Here are some key observations to aid the review process:

🧪 PR contains tests
🔒 No security concerns identified
✅ No TODO sections
🔀 No multiple PR themes
⚡ Recommended focus areas for review

Duplicate Function Signature

Both LOCAL_ARRAY_AGG_OP and LOCAL_ARRAY_AGG_TYPED_OP are registered under the same DataFusion name array_agg in ADDITIONAL_AGGREGATE_SIGS. If the mapping is keyed by name or by the Calcite op, this duplicate registration may cause lookups to be ambiguous or overwrite the earlier entry, potentially breaking the untyped list()/values() path or the new typed path depending on iteration order. Confirm the mapping infrastructure supports multiple Calcite ops mapping to the same backend function name.

 
FunctionMappings.s(LOCAL_ARRAY_AGG_OP, "array_agg"),
FunctionMappings.s(LOCAL_ARRAY_AGG_TYPED_OP, "array_agg"),
Possible Issue

When the input to ARRAY_AGG is already an array type (getComponentType() != null), the rewriter switches to LOCAL_LIST_MERGE_OP with filterArg=-1. However, LIST_MERGE is a FINAL-side merge operator, and using it in place of a PARTIAL/SINGLE ARRAY_AGG for an already-array input may produce incorrect semantics (flattening/un-nesting instead of collecting arrays into an array of arrays). Verify this branch is intended and covered by tests, since mvcombine over an array-typed field is a plausible input.

 
if (arrayAggArg0.getComponentType() != null) {
    arrayAggOp = DataFusionFragmentConvertor.LOCAL_LIST_MERGE_OP;
    arrayAggType = arrayAggArg0;
} else {
    // nullable element: NOT NULL trips Calcite typeMatchesInferred
    arrayAggOp = DataFusionFragmentConvertor.LOCAL_ARRAY_AGG_TYPED_OP;
    RelDataType arrayAggElem = arrayAggTf.createTypeWithNullability(arrayAggArg0, true);
    arrayAggType = arrayAggTf.createTypeWithNullability(arrayAggTf.createArrayType(arrayAggElem, -1), true);
}

@github-actions

github-actions Bot commented Jul 31, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

PR Code Suggestions ✨

Latest suggestions up to f6cfd03

Explore these optional code suggestions:

CategorySuggestion                                                                                                                                    Impact
General
Ensure aggregate result type is nullable

When the input is already an array (component type non-null), the resulting
arrayAggType is assigned directly from the input field type, but the aggregate
result should always be nullable (empty group -> null). Wrap arrayAggArg0 with
createTypeWithNullability(..., true) to avoid a potential NOT NULL array tripping
Calcite's typeMatchesInferred check.

sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/PplAggregateCallRewriter.java [193-201]

 if (arrayAggArg0.getComponentType() != null) {
     arrayAggOp = DataFusionFragmentConvertor.LOCAL_LIST_MERGE_OP;
-    arrayAggType = arrayAggArg0;
+    arrayAggType = arrayAggTf.createTypeWithNullability(arrayAggArg0, true);
 } else {
     // nullable element: NOT NULL trips Calcite typeMatchesInferred
     arrayAggOp = DataFusionFragmentConvertor.LOCAL_ARRAY_AGG_TYPED_OP;
     RelDataType arrayAggElem = arrayAggTf.createTypeWithNullability(arrayAggArg0, true);
     arrayAggType = arrayAggTf.createTypeWithNullability(arrayAggTf.createArrayType(arrayAggElem, -1), true);
 }
Suggestion importance[1-10]: 5

__

Why: Making the aggregate result nullable is a reasonable defensive change since aggregates over empty groups may return null, but it's not clear this is causing an actual bug in the current PR since the input type may already be nullable.

Low
Avoid masking future enum additions

The special-case is placed before the valueOf fallback, but if an ARRAY_AGG enum
constant is ever added later this shortcut will silently mask it. Consider first
attempting the enum lookup and only mapping to LIST when the name isn't a recognized
enum value, or add a comment/assertion documenting the intentional override.

sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AggregateFunction.java [142-146]

-if (name.equalsIgnoreCase("ARRAY_AGG")) {
-    return LIST;
+// mvcombine's ARRAY_AGG uses the LIST decomposition; element type is preserved in
+// PplAggregateCallRewriter (LOCAL_ARRAY_AGG_TYPED_OP). Intentional alias to LIST.
+try {
+    return valueOf(name.toUpperCase(java.util.Locale.ROOT));
+} catch (IllegalArgumentException ignored) {
+    if (name.equalsIgnoreCase("ARRAY_AGG")) {
+        return LIST;
+    }
+    throw new IllegalStateException("Unrecognized aggregate function [" + name + "]");
 }
Suggestion importance[1-10]: 3

__

Why: Minor stylistic concern about future-proofing; the current placement is intentional and the comment already documents the override. Low impact.

Low

Previous suggestions

Suggestions up to commit 7fbfc40
CategorySuggestion                                                                                                                                    Impact
Possible issue
Ensure array return type is nullable

When the input is already an array (component type non-null), the rewrite reuses the
input's array type verbatim as the return type. If the input array is declared NOT
NULL, this will trip Calcite's typeMatchesInferred check on the LIST_MERGE op just
like the scalar branch warns about. Wrap arrayAggArg0 with
createTypeWithNullability(..., true) to force a nullable array return type.

sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/PplAggregateCallRewriter.java [193-196]

 if (arrayAggArg0.getComponentType() != null) {
     arrayAggOp = DataFusionFragmentConvertor.LOCAL_LIST_MERGE_OP;
-    arrayAggType = arrayAggArg0;
+    arrayAggType = arrayAggTf.createTypeWithNullability(arrayAggArg0, true);
 } else {
Suggestion importance[1-10]: 6

__

Why: Reasonable defensive suggestion: if the input array is NOT NULL, reusing it as the return type could trip Calcite's typeMatchesInferred check, similar to the scalar branch. However, it's unclear from the diff whether the input array can be NOT NULL in practice.

Low
General
Handle null name argument safely

Guard against a null name to avoid a NullPointerException from equalsIgnoreCase,
matching the previous behavior where the valueOf call would throw a more specific
error. Consider validating input before the case-insensitive comparison.

sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AggregateFunction.java [141-146]

 public static AggregateFunction fromNameOrError(String name) {
+    if (name == null) {
+        throw new IllegalStateException("Unrecognized aggregate function [null]");
+    }
     // mvcombine's ARRAY_AGG uses the LIST decomposition; element type is preserved in
     // PplAggregateCallRewriter (LOCAL_ARRAY_AGG_TYPED_OP).
     if (name.equalsIgnoreCase("ARRAY_AGG")) {
         return LIST;
     }
Suggestion importance[1-10]: 3

__

Why: Minor defensive null check; the previous code would have thrown NPE on name.toUpperCase() anyway, so this doesn't materially change robustness.

Low
Suggestions up to commit f6cfd03
CategorySuggestion                                                                                                                                    Impact
Possible issue
Ensure aggregate return type is nullable

When the input is already an array (component type not null), reusing the input's
type as the aggregate return type may produce a NOT NULL array that trips Calcite's
typeMatchesInferred check, mirroring the concern noted for the scalar branch. Force
the returned array type to be nullable to keep it consistent with the op's inferred
type.

sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/PplAggregateCallRewriter.java [193-196]

 if (arrayAggArg0.getComponentType() != null) {
     arrayAggOp = DataFusionFragmentConvertor.LOCAL_LIST_MERGE_OP;
-    arrayAggType = arrayAggArg0;
+    arrayAggType = arrayAggTf.createTypeWithNullability(arrayAggArg0, true);
 } else {
Suggestion importance[1-10]: 5

__

Why: The suggestion raises a plausible concern about type nullability consistency with Calcite's typeMatchesInferred check, mirroring the pattern used in the scalar branch. However, without evidence that the input array type is actually NOT NULL in practice, the impact is speculative.

Low
Suggestions up to commit 5dc833e
CategorySuggestion                                                                                                                                    Impact
General
Force nullable array type for merge branch

When the input is already an array (component type != null), the resulting
arrayAggType should still be forced nullable to avoid Calcite's typeMatchesInferred
mismatch, since the merge op's inferred return type is nullable. Wrap arrayAggArg0
with createTypeWithNullability(..., true) in that branch as well.

sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/PplAggregateCallRewriter.java [193-201]

 if (arrayAggArg0.getComponentType() != null) {
     arrayAggOp = DataFusionFragmentConvertor.LOCAL_LIST_MERGE_OP;
-    arrayAggType = arrayAggArg0;
+    arrayAggType = arrayAggTf.createTypeWithNullability(arrayAggArg0, true);
 } else {
     // nullable element: NOT NULL trips Calcite typeMatchesInferred
     arrayAggOp = DataFusionFragmentConvertor.LOCAL_ARRAY_AGG_TYPED_OP;
     RelDataType arrayAggElem = arrayAggTf.createTypeWithNullability(arrayAggArg0, true);
     arrayAggType = arrayAggTf.createTypeWithNullability(arrayAggTf.createArrayType(arrayAggElem, -1), true);
 }
Suggestion importance[1-10]: 5

__

Why: Reasonable defensive change to ensure nullability alignment with the merge op's inferred type, which could prevent Calcite typeMatchesInferred errors, though the actual necessity depends on LOCAL_LIST_MERGE_OP's return type inference.

Low
Avoid NPE on null name argument

Guard against a null name before invoking equalsIgnoreCase, otherwise callers
passing a null aggregate name will hit a NullPointerException here instead of the
previous IllegalStateException from valueOf. Consider validating input first or
reversing the comparison to "ARRAY_AGG".equalsIgnoreCase(name).

sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AggregateFunction.java [141-146]

 public static AggregateFunction fromNameOrError(String name) {
     // mvcombine's ARRAY_AGG uses the LIST decomposition; element type is preserved in
     // PplAggregateCallRewriter (LOCAL_ARRAY_AGG_TYPED_OP).
-    if (name.equalsIgnoreCase("ARRAY_AGG")) {
+    if ("ARRAY_AGG".equalsIgnoreCase(name)) {
         return LIST;
     }
Suggestion importance[1-10]: 3

__

Why: Minor defensive change; swapping to "ARRAY_AGG".equalsIgnoreCase(name) avoids NPE on null input, but it's unclear whether null names are a valid concern here.

Low
Suggestions up to commit 6fc048c
CategorySuggestion                                                                                                                                    Impact
General
Ensure merged array type is nullable

When merging a per-shard array in the FINAL stage, the returned array type should be
nullable to match the aggregate result contract (empty groups → null). Reusing
arrayAggArg0 directly may propagate a NOT NULL array type that trips Calcite's
typeMatchesInferred check downstream. Wrap with createTypeWithNullability(...,
true).

sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/PplAggregateCallRewriter.java [197-200]

 if (arrayAggArg0.getComponentType() != null) {
-    // PARTIAL->FINAL merge: arg0 is already a typed array; preserve it.
+    // PARTIAL->FINAL merge: arg0 is already a typed array; preserve it (nullable array).
     arrayAggOp = DataFusionFragmentConvertor.LOCAL_LIST_MERGE_OP;
-    arrayAggType = arrayAggArg0;
+    arrayAggType = arrayAggTf.createTypeWithNullability(arrayAggArg0, true);
 } else {
Suggestion importance[1-10]: 5

__

Why: Making the merged array type explicitly nullable is a reasonable defensive change to align with Calcite's inferred type expectations, though the input array from a prior aggregate is typically already nullable.

Low
Null-guard the name comparison

Guard against a null name argument before calling equalsIgnoreCase, otherwise a null
input will throw NPE here instead of surfacing the existing "Unrecognized aggregate
function" error path. This preserves the original error semantics for null inputs.

sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AggregateFunction.java [145-147]

-if (name.equalsIgnoreCase("ARRAY_AGG")) {
+if (name != null && name.equalsIgnoreCase("ARRAY_AGG")) {
     return LIST;
 }
Suggestion importance[1-10]: 3

__

Why: Null-guarding preserves original NPE-vs-IllegalStateException semantics, but the original code also would have thrown NPE on name.toUpperCase(...), so this is a minor behavioral consistency improvement.

Low

@github-actions

github-actions Bot commented Jul 31, 2026

Copy link
Copy Markdown
Contributor

✅ Gradle check result for 6fc048c: SUCCESS

@codecov

codecov Bot commented Jul 31, 2026 •
edited
Loading

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 71.40%. Comparing base (03a4de3) to head (f6cfd03).
⚠️ Report is 3 commits behind head on main.

Additional details and impacted files 
@@             Coverage Diff              @@
##               main   #22617      +/-   ##
============================================
- Coverage     71.42%   71.40%   -0.02%     
- Complexity    76786    76789       +3     
============================================
  Files          6148     6148              
  Lines        357980   357980              
  Branches      52177    52177              
============================================
- Hits         255689   255624      -65     
- Misses        81937    82033      +96     
+ Partials      20354    20323      -31

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@noCharger
noCharger force-pushed the mvcombine-datafusion-arrayagg branch from 6fc048c to 5dc833e Compare July 31, 2026 06:43
@github-actions

github-actions Bot commented Jul 31, 2026

Copy link
Copy Markdown
Contributor

Persistent review updated to latest commit 5dc833e

@github-actions

github-actions Bot commented Jul 31, 2026

Copy link
Copy Markdown
Contributor

❌ Gradle check result for 5dc833e: FAILURE

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?

@noCharger
Preserve element type for mvcombine ARRAY_AGG on the DataFusion route
f6cfd03 
PPL mvcombine lowers to Calcite ARRAY_AGG. On the analytics-engine
(DataFusion) route this failed with "No enum constant ...ARRAY_AGG"
(the aggregate was not recognized), and mapping it onto the existing
list()/values() array-agg operator would stringify elements to VARCHAR
(that operator casts to VARCHAR by design for the list()/values()
multivalue-string contract).

Add a dedicated type-preserving path so the collected array keeps the
source element type (ARRAY<INT>, not ARRAY<VARCHAR>):

- AggregateFunction.fromNameOrError: recognize ARRAY_AGG as the LIST
  (collect-into-array) family for backend viability and the
  STATE_EXPANDING decomposition.
- DataFusionFragmentConvertor: add LOCAL_ARRAY_AGG_TYPED_OP whose return
  type is derived from the operand and which does not cast to VARCHAR;
  it drops null elements (matching ARRAY_AGG ... FILTER (IS NOT NULL))
  via filtersNullArgs, and maps to the same type-generic DataFusion
  array_agg UDAF.
- PplAggregateCallRewriter: route ARRAY_AGG to the typed op for the
  PARTIAL/SINGLE stage and to the existing type-generic list_merge for
  the FINAL merge.

list()/values() are unchanged (still ARRAY<VARCHAR>). No Rust change:
the array_agg/list_merge UDAFs already derive the element type.

Signed-off-by: Louis Chu <lingzhichu.clz@gmail.com>
@noCharger
noCharger force-pushed the mvcombine-datafusion-arrayagg branch from 5dc833e to f6cfd03 Compare July 31, 2026 08:59
@github-actions

github-actions Bot commented Jul 31, 2026

Copy link
Copy Markdown
Contributor

Persistent review updated to latest commit f6cfd03

@github-actions

github-actions Bot commented Jul 31, 2026

Copy link
Copy Markdown
Contributor

Persistent review updated to latest commit 7fbfc40

@github-actions

github-actions Bot commented Jul 31, 2026

Copy link
Copy Markdown
Contributor

❌ Gradle check result for 7fbfc40: FAILURE

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?

@noCharger
noCharger force-pushed the mvcombine-datafusion-arrayagg branch from 7fbfc40 to f6cfd03 Compare July 31, 2026 11:18
@github-actions

github-actions Bot commented Jul 31, 2026

Copy link
Copy Markdown
Contributor

Persistent review updated to latest commit f6cfd03

@github-actions

github-actions Bot commented Jul 31, 2026

Copy link
Copy Markdown
Contributor

✅ Gradle check result for f6cfd03: SUCCESS

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@noCharger